package com.nie.token_bucket.Interceptor;

import cn.hutool.json.JSON;
import cn.hutool.json.JSONUtil;
import com.nie.token_bucket.annotation.LimitingRequired;
import com.nie.token_bucket.common.Constant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.ListOperations;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.lang.reflect.Method;
import java.util.Objects;

import static cn.hutool.http.HttpStatus.HTTP_BAD_GATEWAY;

/**
 * 令牌桶拦截器
 *
 * @author nie_cq
 * @date 2022/10/28
 */
@Slf4j
@Component
public class TokenBucketInterceptor extends HandlerInterceptorAdapter {

    @Resource
    private RedisTemplate redisTemplate;

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
        super.postHandle(request, response, handler, modelAndView);
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        HandlerMethod handlerMethod = (HandlerMethod) handler;
        Method method = handlerMethod.getMethod();
        LimitingRequired annotation = method.getAnnotation(LimitingRequired.class);
        if (Objects.nonNull(annotation) && annotation.required()) {
            boolean isAllowAccess = isAllowAccess();
            if (!isAllowAccess) {
                returnJson(response, "当前访问人数过多,请稍后重试!");
                return false;
            }
        }
        return true;
    }

    private void returnJson(HttpServletResponse response, String msg) {
        PrintWriter writer = null;
        response.setCharacterEncoding("utf-8");
        response.setContentType("application/json;charset=utf-8");
        response.setStatus(HTTP_BAD_GATEWAY);
        try {
            writer = response.getWriter();
            ResponseEntity<String> responseEntity = new ResponseEntity<>(msg, HttpStatus.BAD_GATEWAY);
            JSON parse = JSONUtil.parse(responseEntity);
            writer.println(parse);
            writer.flush();
        } catch (IOException e){
            log.error("拦截返回信息错误:"+e.getMessage());
            throw new RuntimeException(e);
        }finally {
            if (null != writer) {
                writer.close();
            }
        }

    }

    /**
     * 是允许访问
     *
     * @return boolean
     */
    private boolean isAllowAccess() {
        ListOperations listOperations = redisTemplate.opsForList();
        Long size = listOperations.size(Constant.TOKEN_BUCKET_KEY);
        if (size < 1) {
            log.info("请求---502");
            return false;
        }
        listOperations.rightPop(Constant.TOKEN_BUCKET_KEY);
        return true;
    }
}
